# -*- coding: utf-8 -*-
"""
Created on Thu Jun  4 22:02:49 2020

@author: Matrix
"""
import torch
import torchvision.models as models
import torch.nn.functional as F
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import cv2
import os
import numpy as np
import threading
from time import *

import PIL.Image as Image
from putlabel import putText

def gesture_recognition(filepath):
    fileList = os.listdir(filepath)
    
    count = 0
    for filename in fileList:
        count += 1
     
    #背景音乐标签
    bgm_label = []
    for i in range(count):
        
        filename = filepath+str(i)+'.jpg'
        #图片读取
        input_image = Image.open(filename)
          
        
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        
        
        #导入测试图片
        input_image = Image.open(filename)
        preprocess = transforms.Compose([
            transforms.Resize(256),
            #transforms.CenterCrop(224),
            transforms.RandomRotation(20),
            #transforms.ColorJitter(contrast=3),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
        input_tensor = preprocess(input_image)
        input_batch = input_tensor.unsqueeze(0) 
        image_tensor = input_batch.to(device)
        
        
        #打开labels
        with open('images/gesture_24.txt', 'r', encoding='gbk') as clf:
            labels = clf.readlines()
            
            
        #导入训练好的模型
        alexnet = torch.load('model/googlenet_model.pkl')
        alexnet.eval()
        
        start = time()
        with torch.no_grad():
            output = alexnet(image_tensor)
 
        prob = F.softmax(output[0], dim=0)
        indexs = torch.argsort(-prob)

        finish = time()
        print("识别时间：")
        print(finish-start)
        #添加音乐标签
        bgm_label.append(labels[indexs[0]].strip())

        #对图片做标记
        putText(filename,labels[indexs[0]])
        
    return bgm_label
       
 
    




#5.0
#定义多线程的类MyThread
class MyThread(threading.Thread):
    
    def __init__(self, func, args, name=''):
        threading.Thread.__init__(self)
        self.name = name
        self.func = func
        self.args = args
        self.result = self.func(*self.args)
 
    def get_result(self):
        try:
            return self.result
        except Exception:
            return None

#多线程图像识别
def gesture_recognition2(filepath):
    

    fileList = os.listdir(filepath)
    
    count = 0
    for filename in fileList:
        count += 1
     
    #初始化背景音乐标签
    bgm_label = []
    for i in range(count):
        bgm_label.append('null')
        
    
    print(count)
    delect = int(count/4)
    print(delect)
    
    #多线程运行识别图像
    threads = []  
    
    t1 = MyThread(ges_rec_thread1, (delect,filepath,), ges_rec_thread1.__name__)
    threads.append(t1)
    t2 = MyThread(ges_rec_thread2, (delect,filepath,), ges_rec_thread2.__name__)
    threads.append(t2)
    t3 = MyThread(ges_rec_thread3, (delect,filepath,), ges_rec_thread3.__name__)
    threads.append(t3)
    t4 = MyThread(ges_rec_thread4, (delect,filepath,count), ges_rec_thread4.__name__)
    threads.append(t4)
      
    for t in threads:
        t.setDaemon(True)
        t.start()
        a = threads[t].get_result()
        print(a)  
    t.join()
    
    print(bgm_label)

def ges_rec_thread1(delect,filepath):
       
    bgm_label1 = []
    for i in range(0,delect):
        
        filename = filepath+str(i)+'.jpg'
        #图片读取
        input_image = Image.open(filename)
          
        
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        
        
        #导入测试图片
        input_image = Image.open(filename)
        preprocess = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
        input_tensor = preprocess(input_image)
        input_batch = input_tensor.unsqueeze(0) 
        image_tensor = input_batch.to(device)
        
        
        #打开labels
        with open('images/gesture_24.txt', 'r', encoding='gbk') as clf:
            labels = clf.readlines()
            
            
        #导入训练好的模型
        alexnet = torch.load('model/googlenet_model.pkl')
        alexnet.eval()

        with torch.no_grad():
            output = alexnet(image_tensor)
 
        prob = F.softmax(output[0], dim=0)
        indexs = torch.argsort(-prob)

        
        #添加音乐标签
        bgm_label1.append(labels[indexs[0]].strip())

        #对图片做标记
        putText(filename,labels[indexs[0]])
        
    return bgm_label1

def ges_rec_thread2(delect,filepath):
    
    bgm_label2 = []
    for i in range(delect,delect*2):
        
        filename = filepath+str(i)+'.jpg'
        #图片读取
        input_image = Image.open(filename)
          
        
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        
        
        #导入测试图片
        input_image = Image.open(filename)
        preprocess = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
        input_tensor = preprocess(input_image)
        input_batch = input_tensor.unsqueeze(0) 
        image_tensor = input_batch.to(device)
        
        
        #打开labels
        with open('images/gesture_24.txt', 'r', encoding='gbk') as clf:
            labels = clf.readlines()
            
            
        #导入训练好的模型
        alexnet = torch.load('model/googlenet_model.pkl')
        alexnet.eval()

        with torch.no_grad():
            output = alexnet(image_tensor)
 
        prob = F.softmax(output[0], dim=0)
        indexs = torch.argsort(-prob)

        
        #添加音乐标签
        bgm_label2.append(labels[indexs[0]].strip())

        #对图片做标记
        putText(filename,labels[indexs[0]])
        
    return bgm_label2

def ges_rec_thread3(delect,filepath):
    
    bgm_label3 = []
    for i in range(delect*2,delect*3):
        
        filename = filepath+str(i)+'.jpg'
        #图片读取
        input_image = Image.open(filename)
          
        
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        
        
        #导入测试图片
        input_image = Image.open(filename)
        preprocess = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
        input_tensor = preprocess(input_image)
        input_batch = input_tensor.unsqueeze(0) 
        image_tensor = input_batch.to(device)
        
        
        #打开labels
        with open('images/gesture_24.txt', 'r', encoding='gbk') as clf:
            labels = clf.readlines()
            
            
        #导入训练好的模型
        alexnet = torch.load('model/googlenet_model.pkl')
        alexnet.eval()

        with torch.no_grad():
            output = alexnet(image_tensor)
 
        prob = F.softmax(output[0], dim=0)
        indexs = torch.argsort(-prob)

        
        #添加音乐标签
        bgm_label3.append(labels[indexs[0]].strip())

        #对图片做标记
        putText(filename,labels[indexs[0]])
        
    return bgm_label3

def ges_rec_thread4(delect,filepath,count):
    
    bgm_label4 = []
    for i in range(delect*3,count):
        
        filename = filepath+str(i)+'.jpg'
        #图片读取
        input_image = Image.open(filename)
          
        
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        
        
        #导入测试图片
        input_image = Image.open(filename)
        preprocess = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
        input_tensor = preprocess(input_image)
        input_batch = input_tensor.unsqueeze(0) 
        image_tensor = input_batch.to(device)
        
        
        #打开labels
        with open('images/gesture_24.txt', 'r', encoding='gbk') as clf:
            labels = clf.readlines()
            
            
        #导入训练好的模型
        alexnet = torch.load('model/googlenet_model.pkl')
        alexnet.eval()

        with torch.no_grad():
            output = alexnet(image_tensor)
 
        prob = F.softmax(output[0], dim=0)
        indexs = torch.argsort(-prob)

        
        #添加音乐标签
        bgm_label4.append(labels[indexs[0]].strip())

        #对图片做标记
        putText(filename,labels[indexs[0]])
        
    return bgm_label4        
        
        

#gesture_recognition2('images/ffmpeg_img/')


















